#include <stdio.h>
#include <getopt.h>
#include <stdlib.h>
#include <string.h>
#include "qmc.h"
#include "params.h"

typedef enum _Args
{
  kArgsWalkers = 0,
  kArgsIterations,
  kArgsWarmup,
  kArgsSeed,
  kArgsReference,
  kArgsDeviceNum,
  kArgsHelp,
  kArgsStep,
  kArgsWaveFunction,
  kArgsWaveParamAlpha,
  kArgsWaveParamNu,
  kArgsWaveParamB,
  kArgsInputFile,
  kArgsOutputFile,
  kArgsWalkersPerGroup
}Args;

void help()
{
  printf("Usage: mchamr [options [opt-arg]] [--input-file filename |"
        "\t--output-file filename grid-x grid-y grid-z]\n");
  printf("options:\n");
  
  printf("\t--opencl-device <device-num>\n");
  printf("\t\tUses the OpenCL device numbered device-num for computation.\n");
  
  printf("\t--help\n");
  printf("\t\tPrints this message. But you already knew that, right?\n");
  
  printf("\t--walkers <num-walkers>\n");
  printf("\t\tThe number of walkers to use in the Monte Carlo simulation.\n");
  
  printf("\t--iterations <num-iterations>\n");
  printf("\t\tThe number of time steps each walker runs.\n");
  
  printf("\t--warmup <iterations>\n");
  printf("\t\tThe number of time steps to warm up the grid before\n");
  printf("\t\taveraging energies. Use with --output-file.\n");
  
  printf("\t--seed <random-seed>\n");
  printf("\t\tThe random seed for SPRNG.\n");
  
  printf("\t--reference\n");
  printf("\t\tDo all calculations using reference CPU implementation.\n");
  
  printf("\t--step <step-size>\n");
  printf("\t\tThe step size of each iteration.\n");
  
  printf("\t--wavefunction <wave>\n");
  printf("\t\tUse 'wave' as Psi. Valid options are 'mcmillan', 'default'.\n");
  
  printf("\t--waveparam-alpha <value>\n");
  printf("\t\tSet alpha parameter in McMillan wave function.\n");
  
  printf("\t--waveparam-nu <value>\n");
  printf("\t\tSet nu parameter in McMillan wave function.\n");
  
  printf("\t--waveparam-b <value>\n");
  printf("\t\tSet b parameter in McMillan wave function.\n");

  printf("\t--input-file <filename>\n");
  printf("\t\tLoad the initial configuration from <filename>. Then do the\n");
  printf("\t\tactual simulation.\n");

  printf("\t--output-file <filename>\n");
  printf("\t\tWrite the initial configuration to <filename> after\n");
  printf("\t\tperforming warmup. Don't so the actual simulation.\n");

  printf("\t--walkers-per-group <group-size>\n");
  printf("\t\tNumber of walkers per walker group. This parameter can\n");
  printf("\t\tspeed up calculations for small numbers of particles on\n");
  printf("\t\tGPUs but shouldn't affect results.\n");

  exit(0);
}

void parseArgs(int argc, char** argv)
{
  static struct option options[] = 
  { 
    {"walkers", required_argument, NULL, kArgsWalkers},
    {"iterations", required_argument, NULL, kArgsIterations},
    {"warmup", required_argument, NULL, kArgsWarmup},
    {"seed", required_argument, NULL, kArgsSeed},
    {"reference", no_argument, NULL, kArgsReference},
    {"opencl-device", required_argument, NULL, kArgsDeviceNum},
    {"help", no_argument, NULL, kArgsHelp},
    {"step", required_argument, NULL, kArgsStep},
    {"wavefunction", required_argument, NULL, kArgsWaveFunction},
    {"waveparam-alpha", required_argument, NULL, kArgsWaveParamAlpha},
    {"waveparam-nu", required_argument, NULL, kArgsWaveParamNu},
    {"waveparam-b", required_argument, NULL, kArgsWaveParamB},
    {"input-file", required_argument, NULL, kArgsInputFile},
    {"output-file", required_argument, NULL, kArgsOutputFile},
    {"walkers-per-group", required_argument, NULL, kArgsWalkersPerGroup},
    {0, 0, NULL, 0}
  };

  //Check options
  while(1)
  {
    int optionIndex = 0;
    int optParsed;

    optParsed = getopt_long(argc, argv, "", options, &optionIndex);
    
    if(optParsed == -1)
    {
      break;
    }

    switch(optParsed)
    {
      case kArgsWalkers:
        if(sscanf(optarg, "%u", &gNumWalkers) != 1)
        {
          printf("--walkers must take an unsigned integer argument.\n");
          exit(1);
        }
        break;
      case kArgsIterations:
        if(sscanf(optarg, "%u", &gNumIterations) != 1)
        {
          printf("--iterations must take an unsigned integer argument.\n");
          exit(1);
        }
        break;
      case kArgsWarmup:
        if(sscanf(optarg, "%u", &gWarmup) != 1)
        {
          printf("--warmup must take an unsigned integer argument.\n");
          exit(1);
        }
        break;
      case kArgsSeed:
        if(sscanf(optarg, "%d", &gRandomSeed) != 1)
        {
          printf("--seed must take an integer argument.\n");
          exit(1);
        }
        break;
      case kArgsReference:
        gUseCPU = true;
        break;
      case kArgsStep:
        if(sscanf(optarg, "%lf", &gStepSize) != 1)
        {
          printf("--step must take a double argument.\n");
          exit(1);
        }
        break;
      case kArgsDeviceNum:
        if(sscanf(optarg, "%u", &gDevNum) != 1)
        {
          printf("--opencl-device must take an integer argument.\n");
          exit(1);
        }
        break;
      case kArgsHelp:
        help();
        break;
      case kArgsWaveFunction:
        if(strcmp(optarg, "default") == 0)
        {
          if(getWaveFunction() != kWaveFunctionNone)
          {
            printf("Error: Only one --wavefuntion may be specified\n");
            exit(1);
          }
          setWaveFunction(kWaveFunctionDefault);
        }
        else if(strcmp(optarg, "mcmillan") == 0)
        {
          if(getWaveFunction() != kWaveFunctionNone)
          {
            printf("Error: Only one --wavefunction may be specified.\n");
            exit(1);
          }
          setWaveFunction(kWaveFunctionMcMillan);
        }
        else
        {
          printf("'%s' unrecognized for --wavefunction argument.\n", optarg);
        }
        break;
      case kArgsWaveParamAlpha:
        if(sscanf(optarg, "%lf", &alphaMcMillan) == 0)
        {
          printf("--waveparam-alpha must take a double precision argument.\n");
          exit(1);
        }
        break;
      case kArgsInputFile:
        gInputFileName = optarg;
        break;
      case kArgsOutputFile:
        gOutputFileName= optarg;
        break;
      case kArgsWalkersPerGroup:
        if(sscanf(optarg, "%u", &gWalkersPerGroup) != 1 
           || gWalkersPerGroup == 0)
        {
          printf("--walkers-per-group must take positive integer argument.\n");
          exit(1);
        }
        break;
      case '?':
        printf("Unrecognized argument: %s\n", optopt);
        exit(1);
        break;
      default:
        printf("Unrecognized argument: %s\n", optopt);
        exit(1);
    }
  }//End check options

  //If we're creating a new initial grid, read its x, y, and z
  if(gOutputFileName != NULL)
  {
    //Make sure we can set x,y,z then do so
    if(optind + 3 > argc)
    {
      printf("When using --output-file to create an initial configuration,\n");
      printf("you must specify the x y z size of the grid. Use --help for\n");
      printf("more info.\n");
      exit(1);
    }

    if(sscanf(argv[optind], "%u", &gXDim) != 1)
    {
      printf("grix-x must be an unsigned integer");
    }
    if(sscanf(argv[optind + 1], "%u", &gYDim) != 1)
    {
      printf("grix-y must be an unsigned integer");
    }
    if(sscanf(argv[optind + 2], "%u", &gZDim) != 1)
    {
      printf("grix-z must be an unsigned integer");
    }

    //Set defaults if not set by options
    if(getWaveFunction() == kWaveFunctionNone)
    {
      setWaveFunction(kWaveFunctionDefault);
    }
  }
  if(gInputFileName != NULL)
  {
    FILE* file;
    unsigned int dummyNumWalkers;

    file = fopen(gInputFileName, "r");
    
    if(file == NULL)
    {
      printf("Error: Could not open input file '%s'.\n", gInputFileName);
      exit(1);
    }

    gYDim = 1;
    gZDim = 1;

    if(fscanf(file, "%u %u\n", &dummyNumWalkers, &gXDim) != 2)
    {
      printf("Format error in %s on line 1.\n", gInputFileName);
      fclose(file);
      exit(1);
    }

    fclose(file);
  }

  derriveParameters();
}

